# -*- coding: utf-8 -*-
"""
混合检索器实现模块

本模块实现了混合检索策略，结合关键词匹配和向量相似度搜索两种方法，
以提供更准确和全面的信息检索能力。

混合检索的优势：
1. 互补性：关键词匹配处理精确匹配，向量搜索处理语义相似性
2. 鲁棒性：当一种方法失效时，另一种方法可以提供备选结果
3. 准确性：通过多种策略的组合，提高检索结果的准确性
4. 灵活性：可以根据不同场景调整各种策略的权重

适用场景：
- 问答系统：结合精确匹配和语义理解
- 文档检索：处理多样化的查询需求
- 知识库搜索：提供多层次的搜索能力
- 智能客服：处理用户的各种表达方式
"""
from difflib import SequenceMatcher
from typing import Any, Dict, List, Optional, Tuple, Union
from loguru import logger
import re
import jieba
from collections import Counter


class HybridRetriever:
    """
    混合检索策略实现类
    
    这个类实现了多种检索策略的组合，包括：
    1. 关键词匹配：基于预定义关键词和字符串相似度
    2. 向量匹配：基于语义向量的相似度搜索
    3. 混合决策：综合多种策略的结果选择最佳答案
    
    设计原则：
    - 策略独立：每种检索策略可以独立工作
    - 权重可调：可以调整不同策略的权重
    - 阈值控制：通过阈值控制结果质量
    - 降级处理：当高质量匹配不可用时提供降级方案
    
    Attributes:
        qa_data (List[Dict]): 问答数据集，包含问题、答案和关键词
        vectorstore: 向量存储实例，用于语义搜索
        keyword_weight (float): 关键词匹配的权重
        similarity_weight (float): 字符串相似度的权重
        vector_threshold (float): 向量匹配的最低阈值
        keyword_threshold (float): 关键词匹配的最低阈值
    
    Examples:
        >>> qa_data = [
        ...     {
        ...         "question": "什么是机器学习？",
        ...         "answer": "机器学习是人工智能的一个分支...",
        ...         "keywords": ["机器学习", "ML", "算法"]
        ...     }
        ... ]
        >>> retriever = HybridRetriever(qa_data, vectorstore)
        >>> answer = retriever.get_best_answer("ML是什么？")
    """
    
    def __init__(self, 
                 qa_data: List[Dict], 
                 vectorstore=None,
                 keyword_weight: float = 0.7,
                 similarity_weight: float = 0.3,
                 vector_threshold: float = 0.3,
                 keyword_threshold: float = 0.5):
        """
        初始化混合检索器
        
        Args:
            qa_data (List[Dict]): 问答数据集，每个元素应包含：
                - question (str): 问题文本
                - answer (str): 答案文本
                - keywords (List[str], optional): 关键词列表
            vectorstore: 向量存储实例，支持similarity_search_with_score方法
            keyword_weight (float): 关键词匹配在总分中的权重 (0.0-1.0)
            similarity_weight (float): 字符串相似度在总分中的权重 (0.0-1.0)
            vector_threshold (float): 向量匹配的最低接受阈值
            keyword_threshold (float): 关键词匹配的最低接受阈值
        
        Raises:
            ValueError: 当权重参数不合理时
            TypeError: 当qa_data格式不正确时
        
        Examples:
            >>> # 基本初始化
            >>> retriever = HybridRetriever(qa_data)
            
            >>> # 自定义权重和阈值
            >>> retriever = HybridRetriever(
            ...     qa_data, 
            ...     vectorstore,
            ...     keyword_weight=0.8,
            ...     similarity_weight=0.2,
            ...     vector_threshold=0.4
            ... )
        """
        # 验证输入参数
        if not isinstance(qa_data, list) or not qa_data:
            raise ValueError("qa_data必须是非空列表")
        
        if not (0 <= keyword_weight <= 1) or not (0 <= similarity_weight <= 1):
            raise ValueError("权重值必须在0-1之间")
        
        if abs(keyword_weight + similarity_weight - 1.0) > 1e-6:
            logger.warning(f"关键词权重({keyword_weight})和相似度权重({similarity_weight})之和不等于1")
        
        # 初始化属性
        self.qa_data = qa_data
        self.vectorstore = vectorstore
        self.keyword_weight = keyword_weight
        self.similarity_weight = similarity_weight
        self.vector_threshold = vector_threshold
        self.keyword_threshold = keyword_threshold
        
        # 验证qa_data格式
        self._validate_qa_data()
        
        # 预处理数据
        self._preprocess_data()
        
        logger.info(f"初始化混合检索器: {len(qa_data)}条数据, 向量存储={'已启用' if vectorstore else '未启用'}")
    
    def _validate_qa_data(self) -> None:
        """
        验证问答数据格式
        
        Raises:
            TypeError: 当数据格式不正确时
        """
        required_fields = ['question', 'answer']
        
        for i, item in enumerate(self.qa_data):
            if not isinstance(item, dict):
                raise TypeError(f"第{i+1}条数据必须是字典类型")
            
            for field in required_fields:
                if field not in item:
                    raise TypeError(f"第{i+1}条数据缺少必需字段: {field}")
                
                if not isinstance(item[field], str) or not item[field].strip():
                    raise TypeError(f"第{i+1}条数据的{field}字段必须是非空字符串")
    
    def _preprocess_data(self) -> None:
        """
        预处理问答数据
        
        为每条数据添加预处理后的信息，如分词结果、标准化文本等。
        """
        for item in self.qa_data:
            # 确保keywords字段存在
            if 'keywords' not in item:
                item['keywords'] = []
            
            # 添加预处理字段
            item['_processed'] = {
                'question_tokens': list(jieba.cut(item['question'])),
                'question_lower': item['question'].lower(),
                'keywords_lower': [kw.lower() for kw in item['keywords']]
            }
    
    def keyword_match(self, question: str) -> Tuple[Optional[Dict], float]:
        """
        基于关键词的匹配
        
        使用关键词匹配和字符串相似度的组合来评估问题的匹配程度。
        
        匹配策略：
        1. 关键词精确匹配：检查预定义关键词是否出现在问题中
        2. 字符串相似度：使用SequenceMatcher计算文本相似度
        3. 加权组合：根据设定的权重组合两种分数
        
        Args:
            question (str): 输入的问题文本
        
        Returns:
            Tuple[Optional[Dict], float]: 最佳匹配的数据项和匹配分数
                - Dict: 匹配的问答数据项，如果没有匹配则为None
                - float: 匹配分数 (0.0-1.0)，分数越高表示匹配度越好
        
        Examples:
            >>> match, score = retriever.keyword_match("什么是机器学习？")
            >>> if score > 0.5:
            ...     print(f"找到匹配: {match['answer']}")
        
        Note:
            关键词匹配对于精确术语和专业词汇特别有效，
            但可能无法处理同义词或语义相似的表达。
        """
        if not question or not question.strip():
            return None, 0.0
        
        question_lower = question.lower()
        question_tokens = set(jieba.cut(question))
        
        best_match = None
        best_score = 0.0
        
        logger.debug(f"开始关键词匹配，问题: {question}")
        
        for item in self.qa_data:
            # 关键词匹配分数
            keyword_score = 0
            keywords = item.get('keywords', [])
            keywords_lower = item['_processed']['keywords_lower']
            
            # 精确关键词匹配
            matched_keywords = 0
            for keyword, keyword_lower in zip(keywords, keywords_lower):
                if keyword_lower in question_lower:
                    matched_keywords += 1
            
            # 计算关键词匹配率
            if keywords:
                keyword_score = matched_keywords / len(keywords)
            
            # 字符串相似度匹配
            similarity = SequenceMatcher(
                None, 
                question_lower, 
                item['_processed']['question_lower']
            ).ratio()
            
            # 加权组合分数
            total_score = (
                keyword_score * self.keyword_weight + 
                similarity * self.similarity_weight
            )
            
            logger.debug(
                f"匹配项: {item['question'][:30]}..., "
                f"关键词分数: {keyword_score:.3f}, "
                f"相似度: {similarity:.3f}, "
                f"总分: {total_score:.3f}"
            )
            
            if total_score > best_score:
                best_score = total_score
                best_match = item
        
        logger.debug(f"关键词匹配完成，最佳分数: {best_score:.3f}")
        return best_match, best_score
    
    def vector_match(self, question: str) -> Tuple[Optional[Dict], float]:
        """
        基于向量的语义匹配
        
        使用预训练的向量模型进行语义相似度搜索，能够理解问题的语义含义
        而不仅仅是字面匹配。
        
        工作原理：
        1. 将问题转换为向量表示
        2. 在向量存储中搜索最相似的向量
        3. 返回对应的问答数据和相似度分数
        
        Args:
            question (str): 输入的问题文本
        
        Returns:
            Tuple[Optional[Dict], float]: 最佳匹配的数据项和相似度分数
                - Dict: 匹配的问答数据项，如果没有匹配则为None
                - float: 相似度分数 (0.0-1.0)，分数越高表示语义越相似
        
        Raises:
            RuntimeError: 当向量搜索过程中出现错误时
        
        Examples:
            >>> match, score = retriever.vector_match("ML的定义是什么？")
            >>> if score > 0.3:
            ...     print(f"语义匹配: {match['answer']}")
        
        Note:
            向量匹配能够处理同义词、语义相似的表达和上下文理解，
            但需要预先构建向量存储，且计算成本相对较高。
        """
        if not self.vectorstore:
            logger.debug("向量存储未初始化，跳过向量匹配")
            return None, 0.0
        
        if not question or not question.strip():
            return None, 0.0
        
        try:
            logger.debug(f"开始向量匹配，问题: {question}")
            
            # 执行向量相似度搜索
            results = self.vectorstore.similarity_search_with_score(question, k=1)
            
            if not results:
                logger.debug("向量搜索无结果")
                return None, 0.0
            
            doc, distance = results[0]
            
            # 将距离转换为相似度分数
            # 不同的向量存储可能使用不同的距离度量
            # 这里假设使用欧几里得距离，转换为0-1的相似度
            similarity = 1.0 / (1.0 + distance)
            
            logger.debug(f"向量匹配结果 - 距离: {distance:.3f}, 相似度: {similarity:.3f}")
            
            # 返回文档的元数据（应包含原始的问答数据）
            return doc.metadata, similarity
            
        except Exception as e:
            error_msg = f"向量检索失败: {e}"
            logger.error(error_msg)
            # 不抛出异常，而是返回空结果，保证系统的鲁棒性
            return None, 0.0
    
    def get_best_answer(self, question: str, return_details: bool = False) -> Union[str, Dict]:
        """
        获取最佳答案
        
        综合关键词匹配和向量匹配的结果，选择最佳的答案。
        
        决策策略：
        1. 优先选择高分的关键词匹配结果（精确匹配）
        2. 如果关键词匹配分数不够高，考虑向量匹配结果
        3. 如果都没有达到阈值，返回最好的可用结果
        4. 如果都没有结果，返回默认消息
        
        Args:
            question (str): 输入的问题文本
            return_details (bool): 是否返回详细的匹配信息
        
        Returns:
            Union[str, Dict]: 
                - 如果return_details=False: 返回答案字符串
                - 如果return_details=True: 返回包含详细信息的字典
        
        Examples:
            >>> # 获取简单答案
            >>> answer = retriever.get_best_answer("什么是AI？")
            >>> print(answer)
            
            >>> # 获取详细信息
            >>> result = retriever.get_best_answer("什么是AI？", return_details=True)
            >>> print(f"答案: {result['answer']}")
            >>> print(f"匹配方法: {result['method']}")
            >>> print(f"置信度: {result['confidence']}")
        """
        if not question or not question.strip():
            error_msg = "问题不能为空"
            return {"answer": error_msg, "method": "error", "confidence": 0.0} if return_details else error_msg
        
        logger.info(f"开始检索问题: {question}")
        
        # 执行关键词匹配
        keyword_match, keyword_score = self.keyword_match(question)
        
        # 执行向量匹配
        vector_match, vector_score = self.vector_match(question)
        
        logger.info(
            f"匹配结果 - 关键词分数: {keyword_score:.3f}, "
            f"向量分数: {vector_score:.3f}"
        )
        
        # 决策逻辑
        selected_match = None
        selected_score = 0.0
        selected_method = "none"
        
        # 优先选择高质量的关键词匹配
        if keyword_score >= self.keyword_threshold:
            selected_match = keyword_match
            selected_score = keyword_score
            selected_method = "keyword"
            logger.info(f"选择关键词匹配结果，分数: {keyword_score:.3f}")
        
        # 如果关键词匹配不够好，考虑向量匹配
        elif vector_score >= self.vector_threshold:
            selected_match = vector_match
            selected_score = vector_score
            selected_method = "vector"
            logger.info(f"选择向量匹配结果，分数: {vector_score:.3f}")
        
        # 降级策略：选择分数最高的结果
        elif keyword_match or vector_match:
            if keyword_score >= vector_score:
                selected_match = keyword_match
                selected_score = keyword_score
                selected_method = "keyword_fallback"
            else:
                selected_match = vector_match
                selected_score = vector_score
                selected_method = "vector_fallback"
            
            logger.warning(
                f"使用降级策略，选择{selected_method}，分数: {selected_score:.3f}"
            )
        
        # 生成结果
        if selected_match:
            answer = selected_match.get('answer', '答案数据异常')
            logger.info(f"找到答案，方法: {selected_method}，置信度: {selected_score:.3f}")
        else:
            answer = "抱歉，我无法找到相关的答案。请尝试重新表述您的问题或联系管理员。"
            logger.warning("未找到匹配的答案")
        
        # 返回结果
        if return_details:
            return {
                "answer": answer,
                "method": selected_method,
                "confidence": selected_score,
                "keyword_score": keyword_score,
                "vector_score": vector_score,
                "matched_data": selected_match
            }
        else:
            return answer
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取检索器统计信息
        
        Returns:
            Dict[str, Any]: 包含统计信息的字典
        
        Examples:
            >>> stats = retriever.get_statistics()
            >>> print(f"数据集大小: {stats['data_size']}")
            >>> print(f"向量存储状态: {stats['vector_enabled']}")
        """
        total_keywords = sum(len(item.get('keywords', [])) for item in self.qa_data)
        avg_keywords = total_keywords / len(self.qa_data) if self.qa_data else 0
        
        return {
            "data_size": len(self.qa_data),
            "vector_enabled": self.vectorstore is not None,
            "total_keywords": total_keywords,
            "avg_keywords_per_item": avg_keywords,
            "keyword_weight": self.keyword_weight,
            "similarity_weight": self.similarity_weight,
            "vector_threshold": self.vector_threshold,
            "keyword_threshold": self.keyword_threshold
        }